import time
import logging
import zmq
import json
from pydiscovery2 import config
from pydiscovery2 import strategy


def gen_rsp_initialize(instrument_id_list: list, is_exception: bool) -> dict:
    assert type(instrument_id_list) == list
    assert type(is_exception) == bool
    rsp = {
        "EVENT_ID": "RSP_INITIALIZE",
        "STRATEGY_ID": config.strategy.strategy_id,
        "TIMESTAMP": time.strftime('%Y-%m-%d %H:%M:%S'),
        "INSTRUMENT_ID_LIST": instrument_id_list,
        "BAR_TIME_LIST": config.strategy.bar_time_list,
        "IS_EXCEPTION": is_exception,
    }
    return rsp


def gen_rsp_push_bar(portfolio: dict, is_replay: bool, is_exception: bool):
    assert type(portfolio) == dict
    assert type(is_replay) == bool
    assert type(is_exception) == bool
    rsp = {
        "EVENT_ID": "RSP_PUSH_BAR",
        "STRATEGY_ID": config.strategy.strategy_id,
        "TIMESTAMP": time.strftime('%Y-%m-%d %H:%M:%S'),
        "PORTFOLIO": portfolio,
        "IS_REPLAY": is_replay,
        "IS_EXCEPTION": is_exception,
    }
    return rsp


def gen_rsp_finalize(is_exception):
    rsp = {
        "EVENT_ID": "RSP_FINALIZE",
        "STRATEGY_ID": config.strategy.strategy_id,
        "TIMESTAMP": time.strftime('%Y-%m-%d %H:%M:%S'),
        "IS_EXCEPTION": is_exception,
    }
    return rsp


def event_loop():
    context = zmq.Context()
    socket = context.socket(zmq.REP)
    socket.bind(config.zmq.addr)

    while True:
        try:
            req_json = socket.recv_string()
            logging.info(f"req_json = {req_json}")
            rsp_json = event_handler(req_json)
            socket.send_string(rsp_json)
            logging.info(f"rsp_json = {rsp_json}")
        except Exception as e:
            logging.error(e)


def event_handler(json_string):
    req_data = json.loads(json_string)
    strategy_id = req_data["STRATEGY_ID"]
    assert strategy_id == config.strategy.strategy_id

    event_id = req_data["EVENT_ID"]
    if event_id == "REQ_INITIALIZE":
        try:
            trading_date = req_data["TRADING_DATE"]
            instrument_id_list = strategy.on_initialize(trading_date)
            rsp_data = gen_rsp_initialize(instrument_id_list=instrument_id_list, is_exception=False)
        except Exception as e:
            rsp_data = gen_rsp_initialize(instrument_id_list=[], is_exception=True)
            logging.error(e)
    elif event_id == "REQ_PUSH_BAR":
        is_replay = req_data["IS_REPLAY"]
        try:
            portfolio = strategy.on_bar(req_data["BAR_FIELD"], req_data["BAR_DATA"])
            rsp_data = gen_rsp_push_bar(portfolio=portfolio, is_replay=is_replay, is_exception=False)
        except Exception as e:
            rsp_data = gen_rsp_push_bar(portfolio={}, is_replay=is_replay, is_exception=True)
            logging.error(e)
    elif event_id == "REQ_FINALIZE":
        try:
            strategy.on_finalize()
            rsp_data = gen_rsp_finalize(is_exception=False)
        except Exception as e:
            rsp_data = gen_rsp_finalize(is_exception=True)
            logging.error(e)
    else:
        raise Exception(f"unexpected event_id({event_id})")
    return json.dumps(rsp_data)
